package com.yihaodian.search.facet;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.lucene.search.CachingCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.grouping.GroupDocs;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.TopGroups;
import org.apache.lucene.search.grouping.term.TermAllGroupsCollector;
import org.apache.lucene.search.grouping.term.TermFirstPassGroupingCollector;
import org.apache.lucene.search.grouping.term.TermSecondPassGroupingCollector;
import org.apache.lucene.util.BytesRef;

import com.yihaodian.api.search.SearchException;
import com.yihaodian.api.search.SearchResponse;
import com.yihaodian.api.search.facet.FacetNode;
import com.yihaodian.api.search.facet.FacetResult;
import com.yihaodian.api.search.facet.SearchFacetSpec;
import com.yihaodian.api.search.grouping.GroupingResult;
import com.yihaodian.api.search.grouping.GroupingSpec;
import com.yihaodian.search.SearchContext;

public class FacetHandler {
	public static void faceting(SearchContext context) throws SearchException{
		try {
			IndexSearcher searcher = context.getSearcher().searcher();

			Map<String, SearchFacetSpec> facetSpecs = context.getShardRequest().
				getSearchRequest().getFacetSpecMap();
			if (facetSpecs == null || facetSpecs.size() == 0) {
				return;
			}
			
			Map<String, FacetResult> resultMap = new HashMap<String, FacetResult>();
			Iterator<Entry<String, SearchFacetSpec>> it = facetSpecs.entrySet().iterator();
			while (it.hasNext()) {
				Entry<String, SearchFacetSpec> entry = it.next();
				GroupingSpec groupingSpec = new GroupingSpec();
				groupingSpec.setGroupingTerm(entry.getKey());
				
				String groupingTerm = groupingSpec.getGroupingTerm();
				Sort groupSort = Sort.RELEVANCE;
				Sort docSort = Sort.RELEVANCE;

				Query query = context.getQuery();

				TermFirstPassGroupingCollector c1 = new TermFirstPassGroupingCollector(
						groupingTerm, groupSort, 10);

				boolean cacheScores = true;
				double maxCacheRAMMB = 4.0;
				CachingCollector cachedCollector = CachingCollector.create(c1,
						cacheScores, maxCacheRAMMB);
				searcher.search(query, context.getFilter(), cachedCollector);

				Collection<SearchGroup<BytesRef>> topGroups = c1.getTopGroups(0,
						true);

				if (topGroups == null) {
					// No groups matched
					return;
				}

				boolean getScores = true;
				boolean getMaxScores = true;
				boolean fillFields = true;
				TermSecondPassGroupingCollector c2 = new TermSecondPassGroupingCollector(
						groupingTerm, topGroups, groupSort, docSort, 10, getScores,
						getMaxScores, fillFields);

				// Optionally compute total group count
				TermAllGroupsCollector allGroupsCollector = null;
				if (true) {
					allGroupsCollector = new TermAllGroupsCollector(groupingTerm);
					// c2 = MultiCollector.wrap(c2, allGroupsCollector);
				}

				if (cachedCollector.isCached()) {
					// Cache fit within maxCacheRAMMB, so we can replay it:
					cachedCollector.replay(c2);
				} else {
					// Cache was too large; must re-execute query:
					searcher.search(query, context.getFilter(), c2);
				}

				TopGroups<BytesRef> groupsResult = c2.getTopGroups(0);
				if (true) {
					groupsResult = new TopGroups<BytesRef>(groupsResult,
							allGroupsCollector.getGroupCount());
				}

				// Generate facetResult...
				FacetResult facetRes = new FacetResult();
				facetRes.nodes = new ArrayList<FacetNode>();
				
				GroupDocs<BytesRef>[] groups = groupsResult.groups;
				for (int i = 0; i < groups.length; i++) {
					GroupingResult result = new GroupingResult();
					GroupDocs<BytesRef> group = groups[i];
					if (group.groupValue != null) {
						FacetNode fn = new FacetNode(group.groupValue.utf8ToString(),
								group.totalHits);
						facetRes.nodes.add(fn);
					}
				}
				Collections.sort(facetRes.nodes);
				resultMap.put(entry.getKey(), facetRes);
			}
			

			SearchResponse response = context.getResponse();
			if (response == null) {
				response = new SearchResponse();
			}
			response.setFacetResultMap(resultMap);
		} catch (Exception e) {
			throw new SearchException(e);
		}
	
	}
}
